package cn.doitedu.spark

import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream}

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DataTypes, StructField, StructType}
import org.roaringbitmap.RoaringBitmap

object Bmxor extends UserDefinedAggregateFunction{
  override def inputSchema: StructType = {
    new StructType(Array(new StructField("bmin",DataTypes.BinaryType)))
  }

  override def bufferSchema: StructType = {
    new StructType(Array(new StructField("buf",DataTypes.BinaryType)))
  }

  override def dataType: DataType = {
   DataTypes.BinaryType
  }

  override def deterministic: Boolean = true

  override def initialize(buffer: MutableAggregationBuffer): Unit = {

    val rbm = new RoaringBitmap()
    val bout = new ByteArrayOutputStream()
    val dout = new DataOutputStream(bout)
    rbm.serialize(dout)

    buffer.update(0,bout.toByteArray)
  }

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val baInput: Array[Byte] = input.getAs[Array[Byte]](0)

    if(buffer.getAs[Array[Byte]](0).size != 0) {

      val baBuff = buffer.getAs[Array[Byte]](0)

      val din = new DataInputStream(new ByteArrayInputStream(baInput))
      val dbf = new DataInputStream(new ByteArrayInputStream(baBuff))

      val rbmIn = new RoaringBitmap()
      val rbmBuf = new RoaringBitmap()

      rbmIn.deserialize(din)
      rbmBuf.deserialize(dbf)

      rbmIn.or(rbmBuf)

      val bout = new ByteArrayOutputStream()
      val dout = new DataOutputStream(bout)
      rbmIn.serialize(dout)

      buffer.update(0, bout.toByteArray)
    }else{
      buffer.update(0,baInput)
    }
  }

  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    val buf1 = buffer1.getAs[Array[Byte]](0)
    val buf2 = buffer2.getAs[Array[Byte]](0)

    val din1 = new DataInputStream(new ByteArrayInputStream(buf1))
    val din2 = new DataInputStream(new ByteArrayInputStream(buf2))

    val rbm1 = new RoaringBitmap()
    val rbm2 = new RoaringBitmap()

    rbm1.deserialize(din1)
    rbm2.deserialize(din2)

    rbm1.or(rbm2)

    val bout = new ByteArrayOutputStream()
    val dout = new DataOutputStream(bout)
    rbm1.serialize(dout)

    buffer1.update(0,bout.toByteArray)

  }

  override def evaluate(buffer: Row): Any = {
    buffer.getAs[Array[Byte]](0)
  }
}
